from maix import camera, display, image, nn, app, uart

# 加载YOLOv5模型
detector = nn.YOLOv5(model="/root/rubbish/best.mud", dual_buff=True)

# 初始化摄像头
cam = camera.Camera(detector.input_width(), detector.input_height(), detector.input_format())

# 初始化显示屏
dis = display.Display()

device = "/dev/ttyS0"

ser = uart.UART(device, 115200)

data = "5"

ser.write(data.encode())
ser.write(data.encode())
ser.write(data.encode())
print(data.encode())

print("初始化成功")

while not app.need_exit():
    # 从摄像头获取图像
    img = cam.read()
    
    # 检测图像中的物体
    objs = detector.detect(img, conf_th=0.80, iou_th=0.45)
    
    # 遍历检测到的所有物体
    for obj in objs:
        # 获取物体的类别名称和置信度
        class_name = detector.labels[obj.class_id]
        cl = obj.class_id
        confidence = obj.score

        if cl == 7 or cl == 8:
            print(1) # 其他
            data = "1"
            ser.write(data.encode())
        elif cl == 2 or cl == 3 or cl == 9:
            print(2) # 可回收
            data = "2"
            ser.write(data.encode())
        elif cl == 4 or cl == 5 or cl == 6:
            print(3) # 厨余垃圾
            data = "3"
            ser.write(data.encode())
        elif cl == 0 or cl == 1:
            print(4) # 有害垃圾
            data = "4"
            ser.write(data.encode())

        # 打印物体的类别和置信度
        print(f"Detected: {class_name}, Confidence: {confidence:.2f}")

        # 绘制矩形框
        img.draw_rect(obj.x, obj.y, obj.w, obj.h, color=image.COLOR_RED)

        # 绘制类别名称和置信度
        msg = f'{class_name}: {confidence:.2f}'
        img.draw_string(obj.x, obj.y, msg, color=image.COLOR_RED)

    # 显示图像
    dis.show(img)
